#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Created on Nov 22,2016
@author by Liuli

"""
import numpy as np
import pandas as pd
import scipy.optimize as sco
import statsmodels.api as sm
import cvxpy as cp
from macro_factor import get_hidden_factor


def factor_exposure(ret, factor_data):
    """ 用 Robust OLS 策略各个资产的beta 暴露"""
    ret.index = pd.to_datetime(ret.index)
    factor_data.index = pd.to_datetime(factor_data.index)

    ret.sort_index(ascending=True, inplace=True)
    factor_data.sort_index(ascending=True, inplace=True)
    factor_data = factor_data.loc[ret.index[0]: ret.index[-1]].copy()

    # 稳健线性回归模型
    exposure = pd.DataFrame(index=factor_data.columns, columns=ret.columns)
    for i in exposure.columns:
        rlm_model = sm.RLM(ret[i], sm.add_constant(factor_data), M=sm.robust.norms.HuberT(), missing='drop').fit()
        exposure[i] = rlm_model.params[exposure.index]

    return exposure


def exposure_contrast(start, end, n=365*2):
    """ 检测不同回溯周期资产相对宏观因子 beta 稳定性情况"""
    start = pd.to_datetime(start).strftime('%Y%m%d')
    _start = (pd.to_datetime(start) - pd.to_timedelta(n+60, 'days')).strftime('%Y%m%d')
    end = pd.to_datetime(end).strftime('%Y%m%d')
    months = pd.date_range(start, end, freq='ME')

    from WindPy import w
    w.start()

    # 300收益, 500收益, 中证1000全收益, 中债-7-10年国债财富(总值)指数, 中债-3-5年期国债财富(总值)指数, 中债-0-1年国债财富(总值)指数,
    # 中债-优选投资级信用债指数, 标普500, SHFE黄金, SHFE铜, INE原油, 货币基金
    index = ['H00300.CSI', 'H00905.CSI', 'H00852.SH', 'CBA06501.CS', 'CBA04601.CS', 'CBA14101.CS', 'CBA20901.CS',
             'SP500TR.SPI', 'AU.SHF', 'CU.SHF', 'SC.INE', 'H11025.CSI']
    asset = w.wsd(index, "close", _start, end, "")
    asset = pd.DataFrame(np.array(asset.Data).T, index=asset.Times, columns=asset.Codes)
    asset.index = pd.to_datetime(asset.index)
    asset = asset[pd.notna(asset['H00300.CSI'])].copy()
    asset = asset.pct_change()

    # 计算高频宏观因子数据
    factor = get_hidden_factor(_start, end)

    exposure = {}
    for m in months:
        start_b = (pd.to_datetime(m) - pd.to_timedelta(n, 'days')).strftime('%Y%m%d')
        exposure[m] = factor_exposure(asset.loc[start_b: m].copy(), factor.loc[start_b: m].copy())

    index = exposure[list(exposure.keys())[0]].index
    columns = exposure[list(exposure.keys())[0]].columns
    weight = {}
    var = pd.DataFrame(index=index, columns=columns)
    for c in columns:
        w = pd.DataFrame(index=index, columns=months)
        for i in months:
            exp = exposure[i][c]
            exp.name = i
            w[i] = exp.loc[w.index]

        weight[c] = w
        var[c] = w.std(axis=1) / w.mean(axis=1).abs()

    return weight, var


def risk_parity(ret, n_trials=20):
    """
    实现风险平价投资组合优化
    
    Parameters:
    -----------
    ret : pd.DataFrame
        资产收益率数据框，每列代表一个资产的收益率序列
        
    Returns:
    --------
    np.array
        最优化后的资产权重向量
    """
    n = ret.shape[1]
    cov_mat = np.array(ret.cov())

    best_result = None
    best_objective = np.inf

    def min_riskparity(weights):
        weights = np.array(weights)
        # 计算组合波动率
        port_vol = np.sqrt(np.dot(weights.T, np.dot(cov_mat, weights)))
        # 计算边际风险贡献
        mrc = np.dot(cov_mat, weights) / port_vol
        # 计算风险贡献
        rc = weights * mrc
        # 计算风险贡献的目标偏差
        rc_target = port_vol / n
        return np.sum((rc - rc_target) ** 2)

    cons = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})
    bnds = tuple((0, 1) for _ in range(n))

    for _ in range(n_trials):
        # 随机初始化权重
        w_init = np.random.dirichlet(np.ones(n))
        
        result = sco.minimize(min_riskparity, w_init,
                            method='SLSQP',
                            bounds=bnds,
                            constraints=cons,
                            options={'ftol': 1e-12, 
                                   'maxiter': 2000,
                                   'disp': False})
        
        if result.fun < best_objective:
            best_objective = result.fun
            best_result = result.x
            
    return best_result


def factor_risk_parity(ret, factor_data, n_trials=20):
    """
    基于宏观因子的风险平价优化函数

    Parameters:
    -----------
    ret : pd.DataFrame
        资产收益率数据
    factor_data : pd.DataFrame
        宏观因子数据，index与ret对齐

    Returns:
    --------
    np.array
        优化后的资产权重向量
    """
    # 准备数据
    n_assets = ret.shape[1]

    # 使用已有的factor_exposure函数计算因子暴露度矩阵
    beta = factor_exposure(ret, factor_data).values  # [n_factors, n_assets]
    factor_cov = factor_data.cov().values  # [n_factors, n_factors]

    best_result = None
    best_objective = np.inf

    def min_factor_riskparity(w):
        w = np.array(w)
        # 将资产权重映射到因子暴露
        factor_weights = beta @ w  # [n_factors]
        # 计算因子层面的组合波动率
        factor_vol = np.sqrt(factor_weights.T @ factor_cov @ factor_weights)
        # 计算因子层面的边际风险贡献
        factor_mrc = factor_cov @ factor_weights / factor_vol
        # 计算因子风险贡献
        factor_rc = factor_weights * factor_mrc
        # 目标：使得各因子的风险贡献相等
        factor_rc_target = factor_vol / len(factor_weights)
        # 最小化风险贡献偏离
        return np.sum((factor_rc - factor_rc_target) ** 2)

    # 约束条件：权重和为1
    cons = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})
    # 约束条件：权重在0到1之间
    bnds = tuple((0, 1) for _ in range(n_assets))

    # 优化求解
    for _ in range(n_trials):
        # 随机初始化权重
        w_init = np.random.dirichlet(np.ones(n_assets))
        
        result = sco.minimize(min_factor_riskparity, w_init,
                            method='SLSQP',
                            bounds=bnds,
                            constraints=cons,
                            options={'ftol': 1e-12, 
                                   'maxiter': 2000,
                                   'disp': False})
        
        if result.fun < best_objective:
            best_objective = result.fun
            best_result = result.x

    return best_result


def check_risk_parity_convergence(weights, ret):
    """检查风险平价优化结果的收敛性"""
    cov_mat = ret.cov().values
    port_vol = np.sqrt(weights @ cov_mat @ weights)
    rc = weights * (cov_mat @ weights) / port_vol
    rc_target = port_vol / len(weights)
    rc_error = np.std(rc) / np.mean(rc)
    print(rc_error)
    return rc_error < 1e-4  # 设定阈值


def risk_parity_scipy(ret, method='trust-constr', n_trials=20):
    """
    改进的风险平价优化函数，使用更高级的优化方法
    
    Parameters:
    -----------
    ret : pd.DataFrame
        资产收益率数据
    method : str
        优化方法，可选 'trust-constr', 'SLSQP', 'COBYLA'
    n_trials : int
        随机初始化次数
        
    Returns:
    --------
    np.array
        最优化后的资产权重向量
    """
    n = ret.shape[1]
    cov_mat = np.array(ret.cov())
    
    def objective(w):
        w = np.array(w)
        port_vol = np.sqrt(w.T @ cov_mat @ w)
        rc = w * (cov_mat @ w) / port_vol
        rc_target = port_vol / n
        return np.sum((rc - rc_target) ** 2)
    
    # 线性约束：权重和为1
    linear_constraint = sco.LinearConstraint(
        np.ones(n),  # 系数矩阵 A
        [1],         # 下界
        [1]          # 上界
    )
    
    # 权重范围约束
    bounds = sco.Bounds(0, 1)
    
    best_result = None
    best_objective = np.inf
    
    for _ in range(n_trials):
        # 使用 Dirichlet 分布生成随机初始权重
        w_init = np.random.dirichlet(np.ones(n))
        
        if method == 'trust-constr':
            # 使用 trust-constr 方法时可以提供 Hessian 矩阵
            result = sco.minimize(
                objective, 
                w_init,
                method='trust-constr',
                constraints=[linear_constraint],
                bounds=bounds,
                options={
                    'verbose': 0,
                    'xtol': 1e-10,
                    'gtol': 1e-10,
                    'barrier_tol': 1e-10,
                    'maxiter': 2000
                }
            )
        else:
            result = sco.minimize(
                objective,
                w_init,
                method=method,
                constraints=[{'type': 'eq', 'fun': lambda x: np.sum(x) - 1}],
                bounds=tuple((0, 1) for _ in range(n)),
                options={
                    'ftol': 1e-12,
                    'maxiter': 2000
                }
            )
        
        if result.success and result.fun < best_objective:
            best_objective = result.fun
            best_result = result.x
    
    # 标准化权重确保精确满足约束
    if best_result is not None:
        best_result = np.clip(best_result, 0, 1)
        best_result = best_result / np.sum(best_result)
    
    return best_result


def factor_risk_parity_scipy(ret, factor_data, method='trust-constr', n_trials=20):
    """
    改进的因子风险平价优化函数
    
    Parameters:
    -----------
    ret : pd.DataFrame
        资产收益率数据
    factor_data : pd.DataFrame
        因子数据
    method : str
        优化方法
    n_trials : int
        随机初始化次数
    """
    n_assets = ret.shape[1]
    beta = factor_exposure(ret, factor_data).values
    factor_cov = factor_data.cov().values
    
    def objective(w):
        w = np.array(w)
        factor_weights = beta @ w
        factor_vol = np.sqrt(factor_weights.T @ factor_cov @ factor_weights)
        factor_mrc = factor_cov @ factor_weights / factor_vol
        factor_rc = factor_weights * factor_mrc
        factor_rc_target = factor_vol / len(factor_weights)
        return np.sum((factor_rc - factor_rc_target) ** 2)
    
    # 线性约束：权重和为1
    linear_constraint = sco.LinearConstraint(
        np.ones(n_assets),
        [1],
        [1]
    )
    
    # 权重范围约束
    bounds = sco.Bounds(0, 1)
    
    best_result = None
    best_objective = np.inf
    
    for _ in range(n_trials):
        w_init = np.random.dirichlet(np.ones(n_assets))
        
        if method == 'trust-constr':
            result = sco.minimize(
                objective,
                w_init,
                method='trust-constr',
                constraints=[linear_constraint],
                bounds=bounds,
                options={
                    'verbose': 0,
                    'xtol': 1e-10,
                    'gtol': 1e-10,
                    'barrier_tol': 1e-10,
                    'maxiter': 2000
                }
            )
        else:
            result = sco.minimize(
                objective,
                w_init,
                method=method,
                constraints=[{'type': 'eq', 'fun': lambda x: np.sum(x) - 1}],
                bounds=tuple((0, 1) for _ in range(n_assets)),
                options={
                    'ftol': 1e-12,
                    'maxiter': 2000
                }
            )
        
        if result.success and result.fun < best_objective:
            best_objective = result.fun
            best_result = result.x
    
    # 标准化权重
    if best_result is not None:
        best_result = np.clip(best_result, 0, 1)
        best_result = best_result / np.sum(best_result)
    
    return best_result


def risk_parity_cvxpy(ret, gamma=100):
    """
    使用 CVXPY 实现风险平价优化
    
    Parameters:
    -----------
    ret : pd.DataFrame
        资产收益率数据
    gamma : float
        风险平价目标的权重系数
        
    Returns:
    --------
    np.array
        最优化后的资产权重向量
    """
    n = ret.shape[1]
    cov_mat = ret.cov().values
    
    # 定义变量
    w = cp.Variable(n)
    risk_contrib = cp.multiply(w, cov_mat @ w)
    
    # 构建目标函数：最小化风险贡献的方差
    target_risk = cp.sum(risk_contrib) / n
    objective = cp.sum_squares(risk_contrib - target_risk)
    
    # 定义约束条件
    constraints = [
        cp.sum(w) == 1,  # 权重和为1
        w >= 0,          # 权重非负
        w <= 1           # 权重上限
    ]
    
    # 构建和求解问题
    prob = cp.Problem(cp.Minimize(objective), constraints)
    prob.solve(solver=cp.MOSEK if 'MOSEK' in cp.installed_solvers() else cp.SCS)
    
    # 返回结果并确保权重和为1
    weights = np.array(w.value)
    weights = np.clip(weights, 0, 1)
    weights = weights / np.sum(weights)
    
    return weights


def factor_risk_parity_cvxpy(ret, factor_data, gamma=100):
    """
    使用 CVXPY 实现基于因子的风险平价优化，采用DCP兼容的方式
    
    Parameters:
    -----------
    ret : pd.DataFrame
        资产收益率数据
    factor_data : pd.DataFrame
        因子数据
    gamma : float
        风险平价目标的权重系数
        
    Returns:
    --------
    np.array
        最优化后的资产权重向量
    """
    n_assets = ret.shape[1]
    n_factors = len(factor_data.columns)
    beta = factor_exposure(ret, factor_data).values
    factor_cov = factor_data.cov().values
    
    # 定义变量
    w = cp.Variable(n_assets)
    risk_budget = cp.Variable(n_factors)  # 风险预算变量
    
    # 计算因子组合风险
    factor_weights = beta @ w
    portfolio_risk = cp.quad_form(factor_weights, factor_cov)
    
    # 构建DCP兼容的目标函数
    # 使用对数障碍函数来实现风险平价
    objective = portfolio_risk
    for i in range(n_factors):
        f_weight = beta[i, :] @ w
        risk_contrib = cp.quad_form(w, factor_cov[i, i] * np.outer(beta[i, :], beta[i, :]))
        objective += gamma * cp.square(cp.log(risk_contrib) - cp.log(risk_budget[i]))
    
    # 定义约束条件
    constraints = [
        cp.sum(w) == 1,          # 权重和为1
        w >= 0.0,                # 权重非负
        w <= 1.0,               # 权重上限
        cp.sum(risk_budget) == 1,  # 风险预算和为1
        risk_budget >= 0.0        # 风险预算非负
    ]
    
    # 构建和求解问题
    prob = cp.Problem(cp.Minimize(objective), constraints)
    
    try:
        # 尝试使用MOSEK求解器（如果可用）
        if 'MOSEK' in cp.installed_solvers():
            prob.solve(solver=cp.MOSEK, verbose=False)
        else:
            # 回退到SCS求解器
            prob.solve(solver=cp.SCS, verbose=False)
            
        if prob.status != cp.OPTIMAL:
            raise ValueError(f"Problem did not converge. Status: {prob.status}")
            
        # 返回结果并确保权重和为1
        weights = np.array(w.value)
        weights = np.clip(weights, 0, 1)
        weights = weights / np.sum(weights)
        
        return weights
        
    except Exception as e:
        print(f"Optimization failed: {str(e)}")
        # 在优化失败时返回等权重组合
        return np.ones(n_assets) / n_assets


if __name__ == "__main__":
    start = '2008-01-01'
    end = '2024-10-31'
    # 备选：中债信用债财富('CBA02701.CS'), 中债-国开行债券总财富(7-10年)指数(CBA02551.CS), ICE布油(B.IPE)
    # 300收益, 500收益, 中证1000全收益, 中债-7-10年国债财富(总值)指数, 中债-3-5年期国债财富(总值)指数, 中债-0-1年国债财富(总值)指数,
    # 中债-优选投资级信用债指数, 标普500, SHFE黄金, SHFE铜, INE原油, 货币基金
    index = ['H00300.CSI', 'H00905.CSI', 'H00852.SH', 'CBA06501.CS', 'CBA04601.CS', 'CBA14101.CS', 'CBA20901.CS',
             'SP500TR.SPI', 'AU.SHF', 'CU.SHF', 'SC.INE', 'H11025.CSI']

    # weight_0_5, var_0_5 = exposure_contrast(start, end, n=182)
    weight_1, var_1 = exposure_contrast(start, end, n=365)
    weight_2, var_2 = exposure_contrast(start, end, n=365*1.5)
    weight_3, var_3 = exposure_contrast(start, end, n=365*2)
    weight_4, var_4 = exposure_contrast(start, end, n=365*3)

    # 保存变量到 Pickle
    import pickle
    f = open('../data/data.pkl', 'wb')
    pickle.dump([weight_1, var_1, weight_2, var_2, weight_3, var_3, weight_4, var_4], f)
    f.close()

    import seaborn as sns
    sns.set_theme(style="whitegrid")
    sns.set_context({'figure.figsize': [12, 6]})
    import matplotlib.pyplot as plt
    plt.rcParams['font.sans-serif'] = 'SimHei'  # windows
    plt.rcParams['axes.unicode_minus'] = False
    pd.set_option('display.max_columns', 20)

    for key in weight_1.keys():
        for factor in weight_1['H00300.CSI'].index:
            weight = pd.DataFrame(index=weight_1[key].columns)
            weight['{}-{}-1'.format(key, factor)] = weight_1[key].loc[factor]
            weight['{}-{}-1.5'.format(key, factor)] = weight_2[key].loc[factor]
            weight['{}-{}-2'.format(key, factor)] = weight_3[key].loc[factor]
            weight['{}-{}-3'.format(key, factor)] = weight_4[key].loc[factor]
            weight.to_excel('./data/report/{}-{}.xlsx'.format(key, factor))
            fig = sns.lineplot(data=weight, palette="tab10", linewidth=2.5)
            line_fig = fig.get_figure()
            line_fig.savefig('./data/report/{}-{}.png'.format(key, factor))
            plt.close()
